{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "88dc24ff",
   "metadata": {},
   "source": [
    "## 用卷积神经网络训练Cifar100\n",
    "For this tutorial, we will use the CIFAR100 dataset. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c2411fc",
   "metadata": {},
   "source": [
    "## 1. Load and normalize CIFAR100\n",
    "Using torchvision, it's extremely easy to load CIFAR100."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7963711",
   "metadata": {},
   "outputs": [],
   "source": [
    "# %matplotlib inline\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import torch.optim.lr_scheduler as lr_scheduler\n",
    "# import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed493286",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the device\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device('cuda')\n",
    "    torch.cuda.set_device(0)\n",
    "else:\n",
    "    device = torch.device('cpu')\n",
    "print('Using PyTorch version:', torch.__version__, ' Device:', device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12bc2962",
   "metadata": {},
   "source": [
    "The output of torchvision datasets are PILImage images of range [0, 1]. We transform them to Tensors of normalized range [-1, 1]."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2dcf4ab3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the batch size and number of epochs\n",
    "batch_size = 128\n",
    "num_epochs = 50\n",
    "\n",
    "# Define the transformation for data preprocessing\n",
    "transform_train = transforms.Compose([\n",
    "            transforms.RandomCrop(32, padding=4),\n",
    "            transforms.RandomHorizontalFlip(),\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize((0.4914, 0.4822, 0.4465),\n",
    "                                 (0.2023, 0.1994, 0.2010)),\n",
    "        ])\n",
    "transform_test = transforms.Compose([\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize((0.4914, 0.4822, 0.4465),\n",
    "                                 (0.2023, 0.1994, 0.2010)),\n",
    "        ])\n",
    "\n",
    "# Load the CIFAR-100 dataset\n",
    "train_dataset = torchvision.datasets.CIFAR100(root='./data', train=True, download=False, transform=transform_train)\n",
    "test_dataset = torchvision.datasets.CIFAR100(root='./data', train=False, download=False, transform=transform_test)\n",
    "\n",
    "# Create data loaders\n",
    "train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2)\n",
    "test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "07f934de",
   "metadata": {},
   "source": [
    "## 2. Define ResNet\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e63816e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "\n",
    "\n",
    "# 定义基本的ResNet块\n",
    "class BasicBlock(nn.Module):\n",
    "    expansion = 1\n",
    "\n",
    "    def __init__(self, in_channels, out_channels, stride=1):\n",
    "        super(BasicBlock, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)\n",
    "        self.bn1 = nn.BatchNorm2d(out_channels)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)\n",
    "        self.bn2 = nn.BatchNorm2d(out_channels)\n",
    "        self.shortcut = nn.Sequential()\n",
    "        if stride != 1 or in_channels != self.expansion * out_channels:\n",
    "            self.shortcut = nn.Sequential(\n",
    "                nn.Conv2d(in_channels, self.expansion * out_channels, kernel_size=1, stride=stride, bias=False),\n",
    "                nn.BatchNorm2d(self.expansion * out_channels)\n",
    "            )\n",
    "\n",
    "    def forward(self, x):\n",
    "        residual = x\n",
    "        out = self.conv1(x)\n",
    "        out = self.bn1(out)\n",
    "        out = self.relu(out)\n",
    "        out = self.conv2(out)\n",
    "        out = self.bn2(out)\n",
    "        out += self.shortcut(residual)\n",
    "        out = self.relu(out)\n",
    "        return out\n",
    "    \n",
    "    \n",
    "class Bottleneck(nn.Module):\n",
    "    expansion = 4\n",
    "\n",
    "    def __init__(self, in_planes, planes, stride=1):\n",
    "        super(Bottleneck, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)\n",
    "        self.bn1 = nn.BatchNorm2d(planes)\n",
    "        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)\n",
    "        self.bn2 = nn.BatchNorm2d(planes)\n",
    "        self.conv3 = nn.Conv2d(planes, self.expansion*planes, kernel_size=1, bias=False)\n",
    "        self.bn3 = nn.BatchNorm2d(self.expansion*planes)\n",
    "\n",
    "        self.shortcut = nn.Sequential()\n",
    "        if stride != 1 or in_planes != self.expansion*planes:\n",
    "            self.shortcut = nn.Sequential(\n",
    "                nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False),\n",
    "                nn.BatchNorm2d(self.expansion*planes)\n",
    "            )\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = torch.relu(self.bn1(self.conv1(x)))\n",
    "        out = torch.relu(self.bn2(self.conv2(out)))\n",
    "        out = self.bn3(self.conv3(out))\n",
    "        out += self.shortcut(x)\n",
    "        out = torch.relu(out)\n",
    "        return out\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c17b5bb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义ResNet模型\n",
    "class ResNet(nn.Module):\n",
    "    def __init__(self, block, layers, num_classes=100):\n",
    "        super(ResNet, self).__init__()\n",
    "        self.in_channels = 64\n",
    "        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)\n",
    "        self.bn1 = nn.BatchNorm2d(64)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n",
    "        self.layer1 = self.make_layer(block, 64, layers[0], stride=1)\n",
    "        self.layer2 = self.make_layer(block, 128, layers[1], stride=2)\n",
    "        self.layer3 = self.make_layer(block, 256, layers[2], stride=2)\n",
    "        self.layer4 = self.make_layer(block, 512, layers[3], stride=2)\n",
    "        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))\n",
    "        self.fc = nn.Linear(512 * block.expansion, num_classes)\n",
    "\n",
    "    def make_layer(self, block, out_channels, blocks, stride=1):\n",
    "        strides = [stride] + [1] * (blocks - 1)\n",
    "        layers = []\n",
    "        for stride in strides:\n",
    "            layers.append(block(self.in_channels, out_channels, stride))\n",
    "            self.in_channels = out_channels * block.expansion\n",
    "\n",
    "        return nn.Sequential(*layers)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.conv1(x)\n",
    "        out = self.bn1(out)\n",
    "        out = self.relu(out)\n",
    "        # out = self.maxpool(out)\n",
    "        out = self.layer1(out)\n",
    "        out = self.layer2(out)\n",
    "        out = self.layer3(out)\n",
    "        out = self.layer4(out)\n",
    "        out = self.avgpool(out)\n",
    "        out = out.view(out.size(0), -1)\n",
    "        out = self.fc(out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2dbf7f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the model\n",
    "model = ResNet(Bottleneck, [3, 4, 6, 3], num_classes=100).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a34acf3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Load the model\n",
    "# PATH = './cifar_skn2.pth'\n",
    "# model.load_state_dict(torch.load(PATH))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12312f6f",
   "metadata": {},
   "source": [
    "## 3. Define a Loss function and optimizer\n",
    "Let’s use a Classification Cross-Entropy loss and Adam."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40fbef67",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the loss function and optimizer\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.SGD(model.parameters(), lr=0.05, momentum=0.9, weight_decay=5e-4)\n",
    "\n",
    "# Define the learning rate schedule\n",
    "lr_schedule_milestones = [int(num_epochs * 0.5), int(num_epochs * 0.75)]  # 50% and 75% of total training epochs\n",
    "lr_schedule_gamma = 0.1  # Learning rate reduction factor\n",
    "\n",
    "# Define the learning rate scheduler\n",
    "scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=lr_schedule_milestones, gamma=lr_schedule_gamma)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ceb4d7b",
   "metadata": {},
   "source": [
    "## 4. Train the network\n",
    "This is when things start to get interesting. We simply have to loop over our data iterator, and feed the inputs to the network and optimize."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "654c8151",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Lists to store training and testing losses, and accuracies\n",
    "train_losses = []\n",
    "test_losses = []\n",
    "train_accs = []\n",
    "test_accs = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f8b32e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training loop\n",
    "total_steps = len(train_loader)\n",
    "for epoch in range(num_epochs):\n",
    "    train_loss = 0.0\n",
    "    train_total = 0\n",
    "    train_correct = 0\n",
    "\n",
    "    for i, (images, labels) in enumerate(train_loader):\n",
    "        images = images.to(device)\n",
    "        labels = labels.to(device)\n",
    "\n",
    "        # Forward pass\n",
    "        outputs = model(images)\n",
    "        loss = criterion(outputs, labels)\n",
    "\n",
    "        # Backward and optimize\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        train_loss += loss.item()\n",
    "        _, predicted = torch.max(outputs.data, 1)\n",
    "        train_total += labels.size(0)\n",
    "        train_correct += (predicted == labels).sum().item()\n",
    "\n",
    "        # Print training progress\n",
    "        if (i + 1) % 100 == 0:\n",
    "            print(\n",
    "                f\"Epoch [{epoch + 1}/{num_epochs}], Step [{i + 1}/{total_steps}], Loss: {train_loss / (i + 1):.4f}\")\n",
    "\n",
    "    train_acc = 100.0 * train_correct / train_total\n",
    "    train_losses.append(train_loss / total_steps)\n",
    "    train_accs.append(train_acc)\n",
    "    \n",
    "    scheduler.step()\n",
    "    # Evaluation on the test set\n",
    "    model.eval()\n",
    "    test_loss = 0.0\n",
    "    test_total = 0\n",
    "    test_correct = 0\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for images, labels in test_loader:\n",
    "            images = images.to(device)\n",
    "            labels = labels.to(device)\n",
    "\n",
    "            # Forward pass\n",
    "            outputs = model(images)\n",
    "            loss = criterion(outputs, labels)\n",
    "\n",
    "            test_loss += loss.item()\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            test_total += labels.size(0)\n",
    "            test_correct += (predicted == labels).sum().item()\n",
    "\n",
    "    test_acc = 100.0 * test_correct / test_total\n",
    "    test_losses.append(test_loss / len(test_loader))\n",
    "    test_accs.append(test_acc)\n",
    "\n",
    "    # Print test accuracy for the current epoch\n",
    "    print(f\"Epoch [{epoch+1}/{num_epochs}], Test Accuracy: {test_acc:.2f}%\")\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da551fa6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Save the model\n",
    "# PATH = './cifar_resnet50_baseline.pth'\n",
    "# torch.save(model.state_dict(), PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c2df4ca",
   "metadata": {},
   "source": [
    "## 5. Test the network on the test data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41c04eae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluation on the test set\n",
    "model.eval()\n",
    "test_loss = 0.0\n",
    "test_total = 0\n",
    "test_correct = 0\n",
    "\n",
    "with torch.no_grad():\n",
    "    for images, labels in test_loader:\n",
    "        images = images.to(device)\n",
    "        labels = labels.to(device)\n",
    "\n",
    "        # Forward pass\n",
    "        outputs = model(images)\n",
    "        loss = criterion(outputs, labels)\n",
    "\n",
    "        test_loss += loss.item()\n",
    "        _, predicted = torch.max(outputs.data, 1)\n",
    "        test_total += labels.size(0)\n",
    "        test_correct += (predicted == labels).sum().item()\n",
    "\n",
    "test_acc = 100.0 * test_correct / test_total\n",
    "test_losses.append(test_loss / len(test_loader))\n",
    "test_accs.append(test_acc)\n",
    "\n",
    "# Print test accuracy for the current epoch\n",
    "print(f\"Test Accuracy of the model on the {len(test_loader.dataset)} test images: {test_acc:.2f}%\")\n",
    "\n",
    "# Switch back to training mode\n",
    "model.train()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6197b53",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the training and testing losses\n",
    "plt.figure(figsize=(10, 5))\n",
    "plt.plot(range(1, num_epochs + 1), train_losses, label='Training Loss')\n",
    "plt.plot(range(1, num_epochs + 1), test_losses, label='Testing Loss')\n",
    "plt.xlabel('Epoch')\n",
    "plt.ylabel('Loss')\n",
    "plt.title('Training and Testing Loss')\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "# Plot the training and testing acc\n",
    "plt.figure(figsize=(10, 5))\n",
    "plt.plot(range(1, num_epochs + 1), train_accs, label='Training Accuracy')\n",
    "plt.plot(range(1, num_epochs + 1), test_accs, label='Testing Accuracy')\n",
    "plt.xlabel('Epoch')\n",
    "plt.ylabel('Accuracy')\n",
    "plt.title('Training and Testing Accuracy')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a810aea9-7d2a-4916-88e4-3c98a92160aa",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
